#!/usr/bin/python2.5
# Copyright (C) 2007 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
  This package is an extension of transitfeed module.
  The goal of this package is to make a gtfs zip file editable.
  We will create temporary files for swapping and
  generate zip files as a result of editing.
"""

import hashlib
import itertools

from datetime import datetime as dt
from jinja2 import Environment, FileSystemLoader

from transitfeed import Schedule, Loader, default_problem_reporter, sqlite, ApproximateDistance

_TEMPLATE_DIR_ = 'res/html'
_JS_DIR_ = 'res/js'

templateEnv = Environment(loader = FileSystemLoader(_TEMPLATE_DIR_))
jsEnv = Environment(loader = FileSystemLoader(_JS_DIR_))


class GTFSBase:
  @staticmethod
  def ToNameDict(gtfsClass, list):
    gtfsobjs = []
    for gtfsobj in list:
      tmp = {}
      idx = 0
      for key in gtfsClass._FIELD_NAMES:
        tmp[key] = gtfsobj[idx]
        if not tmp[key]:
          tmp[key] = ''
        idx += 1
      gtfsobjs.append(tmp)
    return gtfsobjs

  @staticmethod
  def GetAll(gtfsClass, db, filter = '1'):
    list = db.GetAll(gtfsClass._TABLE_NAME, filter)
    return gtfsClass.ToNameDict(list)

  @staticmethod
  def GetPage(gtfsClass, db, filter, page, limit):
    list = db.GetGTFSData(gtfsClass, filter, (page - 1) * limit, limit)
    return gtfsClass.ToNameDict(list)

  @staticmethod
  def GetItemBy(gtfsClass, db, filter = '1'):
    list = db.GetGTFSData(gtfsClass, filter)
    return gtfsClass.ToNameDict(list)[0]
 
  @staticmethod
  def GetItem(gtfsClass, db, id):
    list = db.GetGTFSData(gtfsClass, '`%s` = "%s"' % (gtfsClass._KEY, id))
    return gtfsClass.ToNameDict(list)[0]

  @staticmethod
  def UpdateItemBy(gtfsClass, db, fields, filter):
    db.UpdateGTFSData(gtfsClass, fields, filter)
    db.Commit()

  @staticmethod
  def UpdateItem(gtfsClass, db, fields, id):
    db.UpdateGTFSData(gtfsClass, fields, '`%s` = "%s"' % (gtfsClass._KEY, id))
    db.Commit()

  @staticmethod
  def InsertItem(gtfsClass, db, fields):
    row = []
    for key in gtfsClass._FIELD_NAMES:
      row.append(fields[key])
    db.InsertRows([row], gtfsClass._FIELD_NAMES, gtfsClass._TABLE_NAME)
    db.Commit()


# _TABLE_NAME is the table name of corresponding class in the database.
# The database stores data and intermediate changes which can be used to
# recover editing on the next run of the server and export changes to a
# new GTFS feed.
class Fare:
  _FEED_NAME = 'fare_attributes.txt'
  _TABLE_NAME = 'fares'
  _REQUIRED_FIELD_NAMES = ['fare_id', 'price', 'currency_type',
                           'payment_method', 'transfers']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['transfer_duration']
  _FIELD_TYPES = ['char(256)', 'float', 'char(256)', 'integer',
                  'integer', 'integer']
  _KEY = 'fare_id'

  @staticmethod
  def ToNameDict(list):
    farerules = []
    for farerule in list:
      farerules.append({
                      'id': farerule[0],
                      'price': farerule[1],
                      'currency_type': farerule[2],
                      'payment_method': farerule[3],
                      'transfers': farerule[4],
                      'transfer_duration': farerule[5]
                      })
    return farerules


 
class FareRule:
  _FEED_NAME = 'fare_rules.txt'
  _TABLE_NAME = 'fare_rules'
  _REQUIRED_FIELD_NAMES = ['fare_id']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['route_id',
                                         'origin_id', 'destination_id',
                                         'contains_id']
  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)', 'char(256)',
                  'char(256)']

  @staticmethod
  def ToNameDict(list):
    farerules = []
    for farerule in list:
      farerules.append({
                      'id': farerule[0],
                      'route_id': farerule[1],
                      'origin_id': farerule[2],
                      'destination_id': farerule[3],
                      'contains_id': farerule[4]
                      })
    return farerules

class Transfer:
  _FEED_NAME = 'transfers.txt'
  _TABLE_NAME = 'transfers'
  _REQUIRED_FIELD_NAMES = ['from_stop_id', 'to_stop_id', 'transfer_type']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['min_transfer_time']
  _FIELD_TYPES = ['char(256)', 'char(256)', 'integer', 'integer']

  @staticmethod
  def ToNameDict(list):
    transfers = []
    for transfer in list:
      transfers.append({
                      'from_stop_id': transfer[0],
                      'to_stop_id': transfer[1],
                      'transfer_type': transfer[2],
                      'min_transfer_time': transfer[3]
                      })
    return transfers

"""Utility functions"""
def TimeStringToSeconds(tm):
  if not tm:
    return None
  t = tm.strip().split(':')
  return int(t[0]) * 3600 + int(t[1]) * 60 + int(t[2])

class GTFSDatabase():
  """Initial version of GTFSDatabase, handling of database common
  functions
  TODO(wenxin): Add error handling for all database operation functions
  """

  def __init__(self, file='', memory_db=False):
    if memory_db:
      self._connection = sqlite.connect(':memory:')
    elif file:
      self._connection = sqlite.connect(file + '.db')
    else:
      try:
        self._temp_db_file = tempfile.NamedTemporaryFile()
        self._connection = sqlite.connect(self._temp_db_file.name)
      except sqlite.OperationalError:
        # Windows won't let a file be opened twice. mkstemp does not remove the
        # file when all handles to it are closed.
        self._temp_db_file = None
        (fd, self._temp_db_filename) = tempfile.mkstemp('.db')
        os.close(fd)
        self._connection = sqlite.connect(self._temp_db_filename)

    self._cursor = self._connection.cursor()
    self._file = file

  def Commit(self):
    """Shortcut to connection commit"""
    self._connection.commit()

  def GetAllQuery(self, query):
    self._cursor.execute(query)
    return self._cursor.fetchall()

  def UpdateGTFSData(self, gtfsClass, fields, condition):
    """Get all gtfs data from table"""
    sets = []
    for key in gtfsClass._FIELD_NAMES:
      sets.append('`%s` = "%s"' % (key, fields[key]))

    sql = """update `%s` set %s where %s""" % (
               gtfsClass._TABLE_NAME,
               ','.join(sets),
               condition)
    print sql
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetGTFSData(self, gtfsClass, condition = '1', start = 0, limit = 500):
    """Get all gtfs data from table"""
    sql = """select %s from `%s` where %s limit %s, %s""" % (
               ','.join(gtfsClass._FIELD_NAMES),
               gtfsClass._TABLE_NAME,
               condition,
               start,
               limit)
    print sql
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetAll(self, table_name, condition = 1):
    """Get all from a table"""
    if self.IsTableExist(table_name):
      self._cursor.execute('select * from %s where %s' % (table_name,
        condition))
      return self._cursor.fetchall()
    return []

  def GetStopTimesOfAll(self):
    """Get all stop times info of all trips from table stop_times"""
    sql = ("select trip_id, stop_id, stop_sequence, departure_time,"
           "arrival_time  from stop_times;")
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetStopsOfAll(self):
    """Get all stop info of all trips from table stops"""
    sql = 'select stop_id, stop_lat, stop_lon, stop_name from stops;'
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetTripRouteShapeRelation(self):
    """Get related info between trip, shape and route from table trips"""
    sql = 'select trip_id, shape_id, route_id from trips;'
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def IsGTFSInitialized(self):
    """Check if the database has been initialized correctly"""
    sql = ("select name from sqlite_master where type = 'table' and "
           "name = 'initialized'")
    h = self._cursor.execute(sql)
    ret = h.fetchall()
    return len(ret) == 1

  def SetGTFSAsInitialized(self):
    """Set the Db as correctly initialized by creating
    a table 'initialized' in the Db"""
    sql = 'create table `initialized` (tag text)'
    self._cursor.execute(sql)

  def IsTableExist(self, table_name):
    sql = ("select name from sqlite_master where type = 'table' and"
           " name = '%s'" % (table_name))
    h = self._cursor.execute(sql)
    ret = h.fetchall()
    return len(ret) == 1

  def CreateTable(self, fields, types, name):
    """Create tables for gtfs data."""
    if self._cursor:
      str = 'drop table if exists %s;' % name
      self._cursor.execute(str)
      defs = []
      for i in range(0, len(fields)):
        defs.append(fields[i] + ' ' + types[i])
      str = 'create table %s (%s);' % (name, ','.join(defs))
      self._cursor.execute(str)
      self._connection.commit()

  def CreateIndex(self, fields, name):
    """Create an index for a table.
    
    Args:
      fields: list of field_name needs to create an index such as [field_name] 
              or [field_name1, field_name2]
      name: string of table name
    """
    if self._cursor:
      self._cursor.execute('CREATE INDEX %s_%s_index ON %s (%s);' % 
                           (name, ''.join(fields), name, ','.join(fields)))
      self.Commit()

  def InsertRows(self, rows, header, name):
    """Insert one or more rows from a list."""
    if self._cursor:
      str = ("""INSERT INTO %s (%s) VALUES (%s)""" %  
             (name, ','.join(header), ','.join(['?'] * len(header))))
      self._cursor.executemany(str, rows)

  def AddExtraColumns(self, fields, header, table):
    """Add extra fields not defined in _FIELD_NAMES."""
    for item in header:
      if item not in fields:
        str = 'alter table %s add `%s` char(256);' % (table, item)
        if self._cursor:
          self._cursor.execute(str)

class Calendar:
  """A class inherited from the GenericGTFSObject class
  to enable editing on GTFS calendar.
  """
  _DAYS_OF_WEEK = ['monday', 'tuesday', 'wednesday', 'thursday',
                   'friday', 'saturday', 'sunday']
  _REQUIRED_FIELD_NAMES = ['service_id'] + _DAYS_OF_WEEK + [
                           'start_date', 'end_date']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES
  _TABLE_NAME = 'calendars'
  _FIELD_TYPES = ['char(256)', 'integer', 'integer',
                  'integer', 'integer', 'integer',
                  'integer', 'integer', 'char(256)',
                  'char(256)']
  _FEED_NAME = 'calendar.txt'
  _KEY = 'service_id'

  @staticmethod
  def ToNameDict(list):
    services = []
    for service in list:
      services.append({
                      'id': service[0],
                      'monday': service[1],
                      'tuesday': service[2],
                      'wednesday': service[3],
                      'thursday': service[4],
                      'friday': service[5],
                      'saturday': service[6],
                      'sunday': service[7],
                      'start_date': service[8],
                      'end_date': service[9]
                      })
    services.sort(key = lambda x: x['id'])
    return services

class CalendarDate:
  """A class inherited from the GenericGTFSObject class
  to enable editing on GTFS calendardate.
  """
  _REQUIRED_FIELD_NAMES = ['service_id', 'date', 'exception_type']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + []
  _TABLE_NAME = 'calendar_dates'
  _FIELD_TYPES = ['char(256)', 'char(256)', 'integer']
  _FEED_NAME = 'calendar_dates.txt'
  _KEY = 'service_id'

  @staticmethod
  def ToNameDict(list):
    services = []
    for service in list:
      services.append({
                      'id': service[0],
                      'date': service[1],
                      'exception_type': service[2]
                      })
    return services

class Frequency:
  """A class inherited from the GenericGTFSObject class
  to enable editing on GTFS frequency.
  """
  _REQUIRED_FIELD_NAMES = ['trip_id', 'start_time', 'end_time',
                           'headway_secs']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + []
  _TABLE_NAME = 'frequencies'
  _FIELD_TYPES = ['char(50)', 'char(20)', 'char(20)', 'integer']
  _FEED_NAME = 'frequencies.txt'

  @staticmethod
  def ToNameDict(list):
    frequencies = []
    for frequency in list:
      frequencies.append({
                      'id': frequency[0],
                      'start_time': frequency[1],
                      'end_time': frequency[2],
                      'headway_secs': frequency[3]
                      })
    return frequencies

class Shape:
  _FEED_NAME = 'shapes.txt'
  _TABLE_NAME = 'shapes'
  _REQUIRED_FIELD_NAMES = ['shape_id', 'shape_pt_lat', 'shape_pt_lon',
                           'shape_pt_sequence']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['shape_dist_traveled']
  _FIELD_TYPES = ['char(256)', 'float', 'float', 'integer', 'float']

  @staticmethod
  def ToNameDict(list):
    shapes = []
    for shape in list:
      shapes.append({
                      'id': shape[0],
                      'distance': shape[1]
                      })
    return shapes 

  @staticmethod
  def GetPage(db, filter = '1', page = 1, limit = 5000):
    offset = limit * (page - 1)
    query = ("select shape_id, sum(shape_dist_traveled) as distance from"
             " %s where %s group by shape_id limit %s, %s;" % (
        Shape._TABLE_NAME, filter, offset, limit))
    db._cursor.execute(query)
    rows = db._cursor.fetchall()
    return Shape.ToNameDict(rows)

  @staticmethod
  def UpdateID(db, from_id, to_id):
    query = "update %s set shape_id = '%s' where shape_id = '%s';" % (Shape._TABLE_NAME, to_id, from_id)
    db._cursor.execute(query)

  @staticmethod
  def Exist(db, id):
    query = "select count(*) from shapes where shape_id = '%s';" % id
    db._cursor.execute(query)
    row = db._cursor.fetchone()
    if row and row[0]:
      return 1
    return 0

  @staticmethod
  def SavePoints(db, id, points):
    query = ('delete from %s where shape_id = "%s"' %
            (Shape._TABLE_NAME, id))
    db._cursor.execute(query)
    idx = 1
    rows = []
    for point in points:
      row = point.split(',')
      rows.append([id.encode('utf-8'), row[0], row[1], "%d" % idx, row[2]])
      idx += 1
    db.InsertRows(rows, Shape._FIELD_NAMES, Shape._TABLE_NAME)
    db.Commit()
    return 1

class Stop:
  _FEED_NAME = 'stops.txt'
  _REQUIRED_FIELD_NAMES = ['stop_id', 'stop_name', 'stop_lat', 'stop_lon']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + \
                 ['stop_desc', 'zone_id', 'stop_url', 'stop_code',
                  'location_type', 'parent_station']
  _TABLE_NAME = 'stops'


  _FIELD_TYPES = ['char(256)', 'char(256)', 'float', 'float', 'char(256)',
                  'char(256)', 'char(256)', 'char(256)', 'char(256)',
                  'char(256)']
  _KEY = 'stop_id'
  """Editable Stop inherited from Stop"""

  @staticmethod
  def ToNameDict(list):
    stops = []
    for stop in list:
      stops.append({
                      'id': stop[0],
                      'name': stop[1],
                      'lat': stop[2],
                      'lon': stop[3],
                      'desc': stop[4],
                      'zone_id': stop[5],
                      'url': stop[6],
                      'code': stop[7],
                      'location_type': stop[8],
                      'parent_station': stop[9]
                      })
    stops.sort(key = lambda x: x['name'])
    return stops 

  @staticmethod
  def CopyStop(db, id, parentId, name, lat, lon):
    """Copy the stop info from parentId to create a stop specified by id,
    and update the new stop with the specified name, lat and lon values. """
    query = ("select stop_desc, zone_id, stop_url, stop_code, parent_station"
             " from `%s` where stop_id = '%s'" % 
             (Stop._TABLE_NAME, parentId))
    db._cursor.execute(query)
    row = db._cursor.fetchone()
    # Check if the stop specified by parentId has parent_station.
    # If yes, assign the old parent_station to the new created stop.
    # if not, assign the parentId as the parent_station to the new stop
    if row and row[4]:
      parent_station = row[4]
    else:
      parent_station = parentId 
      Stop.ToStation(db, parentId)
    query = ('insert into `%s` (`stop_id`, `stop_name`, `stop_lat`, '
             '`stop_lon`, `stop_desc`, `zone_id`, `stop_url`, `stop_code`, '
             '`location_type`, `parent_station`) values("%s", "%s", "%s", '
             '"%s","%s", "%s", "%s", "%s", "%s", "%s");' % (
               Stop._TABLE_NAME, id, name, lat, lon, row[0], row[1],
               row[2], row[3], 0, parent_station))
    db._cursor.execute(query)
    query = ("update `%s` set stop_id = '%s' where stop_id = '%s'" %
              (StopTime._TABLE_NAME, id, parentId))
    db._cursor.execute(query)
    db.Commit()
    return parent_station

  @staticmethod
  def GetBoundingBox(db):
    query = ("select min(stop_lat), min(stop_lon), max(stop_lat),"
             "max(stop_lon) from `%s`;" % Stop._TABLE_NAME)

    db._cursor.execute(query)
    return db._cursor.fetchone()

  @staticmethod
  def GetInsertId(db):
    m = hashlib.md5()
    m.update(dt.now().isoformat(' '))
    return m.hexdigest()

  @staticmethod
  def UpdateNameLatLng(db, id, name, lat, lon):
    """Update stop info"""
    query = ('update %s set stop_name = "%s", stop_lat = %s, stop_lon = %s'
           ' where stop_id = "%s"' % (Stop._TABLE_NAME,
                                      name,
                                      lat,
                                      lon,
                                      id))
    db._cursor.execute(query)
    db.Commit()

  @staticmethod
  def ToStation(db, id):
    """Upgrade stop to station by setting it's location type to 1"""
    query = ("update `%s` set location_type = '1' where stop_id = '%s'" %
             (Stop._TABLE_NAME, id))
    db._cursor.execute(query)
    db.Commit()

  @staticmethod
  def Insert(db, id, name, lat, lon, desc='', zone_id='', url='', code='',
      location_type='0', parent_station=''):
    """Insert a new Stop"""
    query = ('insert into `%s` (`stop_id`, `stop_name`, `stop_lat`, '
           '`stop_lon`, `stop_desc`, `zone_id`, `stop_url`, `stop_code`, '
           '`location_type`, `parent_station`) values("%s", "%s", "%s", '
           '"%s","%s", "%s", "%s", "%s", "%s", "%s");' % (
              Stop._TABLE_NAME, id, name, lat, lon, desc, zone_id,
              url, code, location_type, parent_station))
    db._cursor.execute(query)
    db.Commit()

  @staticmethod
  def GetStopsInBoundingBox(db, north, east, south, west, n):
    query = ("select stop_id, stop_name, stop_lat, stop_lon,"
             "parent_station, location_type  from stops "
             "where stop_lat <= '%s' and stop_lat >= '%s'"
             " and stop_lon <= '%s' and stop_lon >= '%s' limit %s" % 
             (north, south, east, west, n))
    return db.GetAllQuery(query)

class StopTime:
  """A class inherited from the GenericGTFSObject class
  to enable editing on GTFS stoptime.
  """
  _TABLE_NAME = 'stop_times'
  _REQUIRED_FIELD_NAMES = ['trip_id', 'arrival_time',
                           'departure_time', 'stop_id',
                           'stop_sequence']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['stop_headsign',
                  'pickup_type', 'drop_off_type',
                  'shape_dist_traveled']
  _FIELD_TYPES = ['char(50)', 'integer', 'integer',
                  'char(50)', 'integer', 'char(100)',
                  'integer', 'integer', 'float']
  _FEED_NAME = 'stop_times.txt'

  @staticmethod
  def ToNameDict(list):
    stoptimes = []
    for stoptime in list:
      tmp = {}
      idx = 0
      for key in StopTime._FIELD_NAMES:
        tmp[key] = stoptime[idx]
        if not tmp[key]:
          tmp[key] = ''
        idx += 1
      stoptimes.append(tmp)
    return stoptimes

  @staticmethod
  def GetStopTimesByTripId(db, tid):
    query = ("select arrival_time, departure_time, stop_sequence,"
             " stops.stop_id, stop_name, stop_lat, stop_lon, parent_station,"
             " location_type from stop_times, stops where trip_id = '%s' and"
             " stops.stop_id = stop_times.stop_id order by stop_sequence;"
             % tid)
    db._cursor.execute(query)
    return db._cursor.fetchall()

  @staticmethod
  def DeleteByPattern(db, routeId, patternId, sequence, stopId):
    """Delete a stoptime record with specified sequence and stop_id
    from trips who share the same pattern_id """
    query = ("delete from `%s` where stop_id = '%s' and stop_sequence = '%s'"
             " and trip_id in (select trip_id from `%s` where"
             " pattern_id = '%s' and route_id = '%s')" % 
             (StopTime._TABLE_NAME, stopId, sequence,
              Pattern._TABLE_NAME, patternId, routeId))
    db._cursor.execute(query)
    db.Commit()

  @staticmethod
  def InsertStopTimeIntoPattern(db, patternId, sequence, arrivalTime,
                                departureTime, stopId, isEnd):
    """Insert stoptime for trips in certain pattern"""
    query = ("select trip_id from `%s` where pattern_id = '%s';" %
      (PatternTrip._TABLE_NAME, patternId))
    db._cursor.execute(query)
    trip_ids = db._cursor.fetchall()

    for trip_id in trip_ids:
      trip_id = trip_id[0]
      # Increase the stop sequence preparing for the newly inserted one
      query = ("update `%s` set stop_sequence = stop_sequence + 1"
               " where stop_sequence >= '%s' and trip_id = '%s';" %
                (StopTime._TABLE_NAME,
                sequence, trip_id))
      db._cursor.execute(query)

      # The following block here is only used to resolve the svg display
      # problem.
      # The svg display procedure needs both the start time and the
      # ending time of the trip to be correctly specified which means they 
      # cannot be empty, so it is necessary to automatically generate such
      # time info for the trip.
      # Now at the frontend javascript code, we automatically copy the 
      # first/last time of the trip to avoid this problem, and update all the
      # trip info of the pattern respectively. 
      # TODO: add the correct departure/arrival time for the inserted stop.

      # BLOCK START
      iTime = ''
      order = ''
      if arrivalTime or departureTime:
        if isEnd:
          order = ' desc '
        db._cursor.execute(
          'select arrival_time, departure_time from stop_times where '
          'trip_id="%s" order by stop_sequence %s limit 1;' % 
          (trip_id, order))
        (arrival_time, departure_time) = db._cursor.fetchone()
        if arrival_time:
          iTime = arrival_time
        elif departure_time:
          iTime = departure_time
      # BLOCK END 

      sql = ('insert into `%s` (trip_id, arrival_time, departure_time,'
             'stop_id, stop_sequence) values("%s", "%s", "%s", "%s", %s);' % (
             StopTime._TABLE_NAME, trip_id, iTime, iTime, stopId,
             sequence))
      db._cursor.execute(sql)
      db.Commit()

class Trip:
  _FEED_NAME = 'trips.txt'
  _TABLE_NAME = 'trips'
  _REQUIRED_FIELD_NAMES = ['route_id', 'service_id', 'trip_id']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['trip_short_name',
    'trip_headsign', 'direction_id', 'block_id', 'shape_id'
    ]
  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)', 
                  'char(256)', 'char(256)', 'integer',
                  'char(256)', 'char(256)']
  _KEY = 'trip_id'

  """A class inherited from the Trip class to enable editing
  on GTFS feeds.
  """

  @staticmethod
  def ToNameDict(list):
    trips = []
    for trip in list:
      tmp = {}
      idx = 0
      for key in Trip._FIELD_NAMES:
        tmp[key] = trip[idx]
        if not tmp[key]:
          tmp[key] = ''
        idx += 1
      trips.append(tmp)
    return trips

  @staticmethod
  def GetShapePointsByTripId(db, tid):
    query = ("select shape_pt_lat, shape_pt_lon, shape_dist_traveled"
             " from shapes, trips where trips.trip_id = '%s' and"
             " shapes.shape_id = trips.shape_id;" % tid)
    return db.GetAllQuery(query)

  @staticmethod
  def GetPatternTripsByRouteId(db, rid):
    query = ("select patterntrips.pattern_id, trips.trip_id,"
             " trips.trip_headsign, trips.service_id, trips.direction_id,"
             " trips.block_id, trips.shape_id, patterns.pattern_desc,"
             " patterntrips.start_time from patterntrips, trips, patterns"
             " where trips.route_id = '%s' and"
             " patterns.route_id = trips.route_id and"
             " patterntrips.trip_id = trips.trip_id;" % rid)
    return db.GetAllQuery(query)

  @staticmethod
  def GetStopLatLngsByTripId(db, trip_id):
    """Get stop latlngs of the trip specified by the trip_id"""
    query = ("select stop_lat, stop_lon from stop_times, stops"
             " where trip_id = '%s' and stops.stop_id = stop_times.stop_id"
             " order by stop_sequence;"
             % trip_id)
    db._cursor.execute(query)
    return db._cursor.fetchall()

  @staticmethod
  def GetStopGapsByTripId(db, trip_id):
    """Get the stop distance between the nearest stops in the trip"""
    latlngs = Trip.GetStopLatLngsByTripId(db, trip_id)

    gaps = [ApproximateDistance(stop1[0], stop1[1], stop2[0], stop2[1]) for
        (stop1, stop2) in itertools.izip(latlngs, latlngs[1:])]
    return gaps

  @staticmethod
  def GetStopTimesByTripId(db, trip_id):
    """Get the time list of the trip"""
    query = ("select arrival_time, departure_time from stop_times"
             " where trip_id = '%s' order by stop_sequence" % (trip_id))
    db._cursor.execute(query)
    return db._cursor.fetchall()

class EditableSchedule(Schedule):
  """A class inherited from the Schedule class to enable editing
  on GTFS feeds.
  """

  def __init__(self, feed_path=None,
               problem_reporter=default_problem_reporter,
               memory_db=False, check_duplicate_trips=False):
    """an overriding on the parent __init__ function.
    add _original_shapes"""
    self._feed_path = feed_path
    Schedule.__init__(self, problem_reporter, memory_db,
                      check_duplicate_trips)
    self._original_shapes = {}  # Original shape_id to Shape

  def SaveToFeedFile(self):
    """Save database info into a zip file"""
    basename, extension = os.path.splitext(self._feed_path)
    self.WriteGoogleTransitFeed(basename + "-new" + extension)    

  def ExportGTFSDataToTextFile(self, gtfsClass, archive):
    """Export a gtfsClass data into an archive"""
    str = StringIO.StringIO()
    writer = CsvUnicodeWriter(str)
    columns = gtfsClass._FIELD_NAMES
    writer.writerow(columns)
    writer.writerows(self._db.GetGTFSData(gtfsClass))
    self._WriteArchiveString(archive, gtfsClass._FEED_NAME, str)

  def WriteGoogleTransitFeed(self, file):
    """Output this schedule as a Google Transit Feed Data in file """
    archive = zipfile.ZipFile(file, 'w')
    classes = [
      Agency, Stop, Route,
      Calendar, CalendarDate, StopTime,
      Trip, Shape, Frequency,
      Fare, FareRule, Transfer]
    for gtfsClass in classes:
      self.ExportGTFSDataToTextFile(gtfsClass, archive)
    archive.close()

  def ConnectDb(self, memory_db):
    """Overridden to enable path specified db"""
    self._db = GTFSDatabase(self._feed_path, memory_db=memory_db)
    self._connection = self._db._connection
    self._cursor = self._db._connection.cursor()

  def Load(self, feed_path, extra_validation=False):
    """An overriding on the parent Load. Check the import
    and add flag for it."""
    self._imported = self._db.IsGTFSInitialized()

    loader = EditableLoader(feed_path,
                            self,
                            problems=self.problem_reporter,
                            extra_validation=extra_validation,
                            imported=self._imported,
                            connection=self._connection,
                            db=self._db)
    loader.Load()

    if not self._imported:
      self._db.SetGTFSAsInitialized()

class EditableLoader(Loader):
  """Inherited from the Loader to enable editing on GTFS feeds."""
  def __init__(self,
               feed_path=None,
               schedule=None,
               problems=default_problem_reporter,
               extra_validation=False,
               load_stop_times=True,
               memory_db=True,
               zip=None,
               check_duplicate_trips=False,
               imported=False,
               connection=None,
               db=None):
    self._imported = imported
    if connection:
      self._connection = connection
      self._db = db
      self._cursor = self._db._cursor

    Loader.__init__(self, feed_path, schedule, problems, extra_validation, 
                    load_stop_times, memory_db, zip, check_duplicate_trips)

  def _ImportDataFromCsvToDb(self, gtfs_obj):
    """Import data from a file in gtfs feed into a table in database
    Args:
      gtfs_obj: a class drived from GenericGTFSObject
    """
    extra_columns_added = False
    self._db.CreateTable(gtfs_obj._FIELD_NAMES,
                         gtfs_obj._FIELD_TYPES,
                         gtfs_obj._TABLE_NAME)

    if not self._HasFile(gtfs_obj._FEED_NAME):
      return

    rows = []
    header = []
    for (d, row_num, header, row) in \
     self._ReadCsvDict(gtfs_obj._FEED_NAME,
                       gtfs_obj._FIELD_NAMES,
                       gtfs_obj._REQUIRED_FIELD_NAMES):
      if not extra_columns_added:
        columns = self._db.AddExtraColumns(gtfs_obj._FIELD_NAMES,
                                           header,
                                           gtfs_obj._TABLE_NAME)
        extra_columns_added = True
      rows.append(row)

    # Store all the listings in rows and insert into DB with executemany(), 
    # this can reduce 1/3 time.
    if rows:
      self._db.InsertRows(rows, header, gtfs_obj._TABLE_NAME)
    del rows

  def _LoadObject(self, gtfs_obj):
    """Always initialize from database, and import data into
    database if it is not imported"""
    self._ImportDataFromCsvToDb(gtfs_obj)

  def _BindTripShapeStopTimesInfo(self):
    """initialize relationship between entities"""

    rows = []
    pattern_rows = []
    # Self.GetStopTimesOfAll() returns an array in format of
    # (trip_id, stop_id, stop_sequence, departure_time, arrival_time)
    stoptimes_of_all = self._db.GetStopTimesOfAll()
    stoptimes_in_trip = {}
    for stoptime in stoptimes_of_all:
      stoptimes_in_trip.setdefault(stoptime[0], []).append([stoptime[1],
                                                            stoptime[2],
                                                            stoptime[3],
                                                            stoptime[4]])
    for trip_id in stoptimes_in_trip:
      def sort_stoptimes(t1, t2):
        return cmp(t1[1], t2[1])
      stoptimes_in_trip[trip_id].sort(sort_stoptimes)

    # Self._db.GetStopsOfAll() returns an array in format of
    # (stop_id, stop_lat, stop_lon, stop_name)
    stops_of_all = self._db.GetStopsOfAll()
    stop_dict = {}
    for stop in stops_of_all:
      stop_dict.setdefault(stop[0], []).append(stop)

    hashed_shapes = {}
    full_info_stop_lists = {}

    # Self._db.GetTripRouteShapeRelation() returns an array in format of
    # (trip_id, shape_id, route_id)
    trip_shape_route_relation = self._db.GetTripRouteShapeRelation()
    for relation in trip_shape_route_relation:
      # Generate pattern_id with the hash value of a tuple of stop ids.
      # Initialize the relationships and store them into tables.
      # It will cost some time. The pattern_id will be generated once
      # before the feed initialized.
      trip_id = relation[0]
      shape_id = relation[1]
      route_id = relation[2]
      stoptimes = stoptimes_in_trip[trip_id]
      stopids = ["%s" % v[0] for v in stoptimes]
      pattern_id = hash(tuple(stopids))
      str = ''.join(stopids)
      m = hashlib.md5()
      m.update(str + route_id)
      shape_hash = m.hexdigest()
      if not shape_hash in hashed_shapes:
        hashed_shapes[shape_hash] = False

      if not hashed_shapes[shape_hash]:
        full_stop_info_list = []
        for stoptime in stoptimes:
          stop = stop_dict[stoptime[0]][0] # Get first stop
          full_stop_info = [stoptime[1],
                            stop[1],
                            stop[2],
                            stop[3],
                            stoptime[3]]
          full_stop_info_list.append(full_stop_info)
        full_info_stop_lists[shape_hash] = full_stop_info_list
        hashed_shapes[shape_hash] = True
      if not shape_id:
        shape_id = shape_hash
        self._db._cursor.execute('UPDATE trips set shape_id = "%s" WHERE '
                           'trip_id = "%s";' % 
                           (shape_id, trip_id))
      else:
        shape_id_hashed = False
      full_stop_info_list = full_info_stop_lists[shape_hash]
      shape_desc = '%s to %s, %d stops' % (full_stop_info_list[0][3],
                                           full_stop_info_list[-1][3],
                                           len(full_stop_info_list))

      start_time = stoptimes[0][3]
      rows.append([pattern_id, trip_id, start_time])
      if not (route_id, pattern_id) in [(row[0], row[1]) \
          for row in pattern_rows]:
        pattern_rows.append([route_id, pattern_id, shape_desc])
    self._db.InsertRows(rows, PatternTrip._FIELD_NAMES,
      PatternTrip._TABLE_NAME)
    self._db.InsertRows(pattern_rows, Pattern._FIELD_NAMES,
      Pattern._TABLE_NAME)

  def InitIndex(self):
    """Create indexes on classes to improve the query performance later"""
    indexRows = [
      ['agency_id', Agency._TABLE_NAME],
      ['service_id', CalendarDate._TABLE_NAME],
      ['service_id', Calendar._TABLE_NAME],
      ['fare_id', Fare._TABLE_NAME],
      ['fare_id', FareRule._TABLE_NAME],
      ['route_id', FareRule._TABLE_NAME],
      ['pattern_id', Pattern._TABLE_NAME],
      ['route_id', Pattern._TABLE_NAME],
      ['pattern_id', PatternTrip._TABLE_NAME],
      ['trip_id', PatternTrip._TABLE_NAME],
      ['route_id', Route._TABLE_NAME],
      ['agency_id', Route._TABLE_NAME],
      ['trip_id', Frequency._TABLE_NAME],
      ['trip_id', Trip._TABLE_NAME],
      ['route_id', Trip._TABLE_NAME],
      ['service_id', Trip._TABLE_NAME],
      ['shape_id', Trip._TABLE_NAME],
      ['shape_id', Shape._TABLE_NAME],
      ['stop_id', Stop._TABLE_NAME],
      ['zone_id', Stop._TABLE_NAME],
      ['trip_id', StopTime._TABLE_NAME],
      ['stop_id', StopTime._TABLE_NAME],
      ['from_stop_id', Transfer._TABLE_NAME],
      ['to_stop_id', Transfer._TABLE_NAME]]

    for row in indexRows:
      (index, table) = row
      self._db.CreateIndex([index], table)

  def Load(self):
    self._problems.ClearContext()
    if not self._DetermineFormat():
      return self._schedule
    
    self._CheckFileNames()

    # TODO: Add validations for all the classes of GTFS
    if not self._imported:
      gtfsClasses = [Agency, Stop, Route, Calendar, CalendarDate,
          Transfer, Trip, Shape, Frequency,
          Fare, FareRule, StopTime]
      for gtfs in gtfsClasses:
        self._LoadObject(gtfs)
     
      # Init Pattern
      self._db.CreateTable(Pattern._FIELD_NAMES,
                           Pattern._FIELD_TYPES,
                           Pattern._TABLE_NAME)
      self._db.CreateTable(PatternTrip._FIELD_NAMES,
                           PatternTrip._FIELD_TYPES,
                           PatternTrip._TABLE_NAME)
      self.InitIndex()
      self._BindTripShapeStopTimesInfo()

    if self._zip:
      self._zip.close()
      self._zip = None

    return self._schedule

class PatternTrip():
  _TABLE_NAME = 'patterntrips'
  _FIELD_NAMES = ['pattern_id', 'trip_id', 'start_time']
  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)']

  @staticmethod
  def getPatternIdByTripId(db, trip_id):
    query = ("select pattern_id from patterntrips where trip_id = '%s'" %
             trip_id)
    db._cursor.execute(query)
    return db._cursor.fetchone()[0]

  @staticmethod
  def getPatternMinMaxTime(db, pattern_id):
    """Get the min and max time of the pattern defined in the stop_times
    by the trips in the pattern"""

    # Currently arrival_time and departure_time are of differences,
    # so we are now randomly chosed one for min/max
    query = ("select min(a.arrival_time), max(a.departure_time) from"
             " stop_times as a, trips as b, patterntrips as c where"
             " a.arrival_time != '' and a.trip_id = b.trip_id"
             " and b.trip_id = c.trip_id and c.pattern_id = '%s'" %
             (pattern_id))
    db._cursor.execute(query)
    return db._cursor.fetchone()

  @staticmethod
  def getTripsByPatternId(db, pattern_id):
    """Get the trips of the pattern"""
    query = ("select a.trip_id, service_id, start_time"
             " from patterntrips as a, trips as b"
             " where a.trip_id = b.trip_id and pattern_id = '%s'" %
             pattern_id)
    db._cursor.execute(query)
    return db._cursor.fetchall()

class Pattern():
  _TABLE_NAME = 'patterns'
  _FIELD_NAMES = ['route_id', 'pattern_id', 'pattern_desc']
  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)']

class Agency:
  _FEED_NAME = 'agency.txt'
  _TABLE_NAME = 'agencies'
  _REQUIRED_FIELD_NAMES = ['agency_name', 'agency_url', 'agency_timezone']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['agency_id', 'agency_lang',
                                          'agency_phone']

  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                  'char(256)', 'char(256)', 'char(256)']
  _KEY = 'agency_id'

  @staticmethod
  def ToNameDict(agenciesList):
    agencies = []
    for agency in agenciesList:
      agencies.append({
                      'id': agency[3],
                      'name': agency[0],
                      'url': agency[1],
                      'timezone': agency[2],
                      'lang': agency[4],
                      'phone': agency[5]
                      })
    agencies.sort(key = lambda x: x['name'])
    return agencies

class Route:
  _FEED_NAME = 'routes.txt'
  _TABLE_NAME = 'routes'
  _REQUIRED_FIELD_NAMES = [
    'route_id', 'route_short_name', 'route_long_name', 'route_type'
    ]
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + [
    'agency_id', 'route_desc', 'route_url', 'route_color', 'route_text_color'
    ]

  _FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                  'integer', 'char(256)', 'char(256)',
                  'char(256)', 'char(256)', 'char(256)']
  _KEY = 'route_id'

  @staticmethod
  def ToNameDict(routesList):
    routes = []
    for route in routesList:
      routes.append({
                    'id': route[0],
                    'shortName': route[1],
                    'longName': route[2],
                    'type': route[3],
                    'agencyId': route[4],
                    'desc': route[5],
                    'url': route[6],
                    'color': route[7],
                    'textColor': route[8]
                    })
    def sort_routes(r1, r2):
      rn1 = r1['longName'] or r1['shortName']
      rn2 = r2['longName'] or r2['shortName']
      return cmp(rn1, rn2)
    routes.sort(sort_routes)
    return routes
